import os
import glob
import argparse
from gmssl import *

def append_to_file(file_path, content):
    """追加内容到文本文件末尾"""
    with open(file_path, 'a', encoding='utf-8') as f:  # 'a' 表示追加模式
        f.write(content)
        #f.write(content + '\n')  # 自动添加换行符

def safe_process_large_file(input_file, output_file, chunk_size=1024*1024):
    try:
        with open(input_file, 'rb') as f:
            chunk_count = 0
            while True:
                chunk = f.read(chunk_size)
                if not chunk:
                    break
                
                # 处理当前块
                append_to_file(output_file, chunk.hex().upper())
                chunk_count += 1
                
        print(f"文件处理完成,共 {chunk_count} 块\n")
        
    except MemoryError:
        print("内存不足，尝试减小块大小")
    except Exception as e:
        print(f"处理失败: {e}")

def read_binary_file(file_path):
    """读取二进制文件内容并转换为字符串"""
    with open(file_path, 'rb') as f:
        binary_data = f.read()
        # 尝试解码为字符串，如果失败则用十六进制表示
        return binary_data.hex().upper()

def hash(msg_file, dgst_file):
    with open(msg_file, 'rb') as f:
        msg_data = f.read()

    with open(dgst_file, 'rb') as f:
        dgst_data = f.read()

    sm3 = Sm3()
    sm3.update(msg_data)
    dgst = sm3.digest()
    assert dgst == dgst_data, "杂凑值错误"

def main():
    parser = argparse.ArgumentParser(description='遍历目录中的二进制文件')
    parser.add_argument('directory', help='目标目录路径')
    args = parser.parse_args()
    directory = args.directory

    if not os.path.isdir(directory):
        print(f"错误：目录不存在 - {directory}")
        return
    
    # 创建输出内容
    merged_content = ""
    
    msg_file = os.path.join(directory, f'msg')
    hash_file = os.path.join(directory, f'hash')
    hash (msg_file, hash_file) 

    # 读取文件内容
    merged_content = f"消息=\n"
    output_file = 'sm3速率.txt'

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(merged_content)

    safe_process_large_file(msg_file, output_file, chunk_size=0x20000)

    hash_content = read_binary_file(hash_file)
    merged_content = f"\n杂凑值=\n{hash_content}\n"
    with open(output_file, 'a', encoding='utf-8') as f:
        f.write(merged_content)


    print(f"合并完成，结果已保存到 {output_file}")

if __name__ == "__main__":
    main()
